# 使用 openwebui 的环境运行

# pip install requests python-dotenv tqdm
# pip install retry python-dotenv tqdm

import os
import requests
from tqdm import tqdm
from dotenv import load_dotenv

load_dotenv()

class OpenWebUIBatchUploader:
    def __init__(self):
        self.base_url = os.getenv("OPENWEBUI_URL", "http://localhost:8080")
        self.api_key = os.getenv("OPENWEBUI_API_KEY", "sk-9d9eec74afb04156a19be4f14d9613b0")
        self.kb_id = os.getenv("KNOWLEDGE_BASE_ID", "20536a96-5b3c-43b0-9b53-b590caae2c60")  # 使用知识库UUID
        self.headers = {
            "Authorization": f"Bearer {self.api_key}",
            "accept": "application/json"
        }

    def _upload_file(self, file_path):
        """第一步：上传文件到服务器获取file_id"""
        url = f"{self.base_url}/api/v1/files/"

        try:
            with open(file_path, "rb") as f:
                files = {"file": (os.path.basename(file_path), f)}
                response = requests.post(url, files=files, headers=self.headers)

            if response.status_code == 200:
                return True, response.json()["id"]
            return False, response.text
        except Exception as e:
            return False, str(e)

    def _add_to_knowledge(self, file_id):
        """第二步：将文件添加到知识库"""
        url = f"{self.base_url}/api/v1/knowledge/{self.kb_id}/file/add"

        payload = {"file_id": file_id}
        try:
            response = requests.post(url, json=payload, headers=self.headers)
            return response.status_code == 200, response.text
        except Exception as e:
            return False, str(e)

    def process_file(self, file_path):
        """处理单个文件的全流程"""
        # 第一步上传文件
        upload_status, upload_result = self._upload_file(file_path)
        if not upload_status:
            return False, f"Upload failed: {upload_result}"

        # 第二步添加到知识库
        add_status, add_result = self._add_to_knowledge(upload_result)
        if not add_status:
            return False, f"Add to KB failed: {add_result}"

        return True, "Success"

    def batch_upload(self, directory_path):
        """批量处理目录中的文件"""
        supported_ext = ['.pdf', '.txt', '.docx', '.md', '.csv', '.html']
        files = [
            f for f in os.listdir(directory_path)
            if os.path.splitext(f)[1].lower() in supported_ext
        ]

        if not files:
            print("目录中没有支持的文件格式")
            return

        success = 0
        failures = []

        for filename in tqdm(files, desc="上传进度"):
            file_path = os.path.join(directory_path, filename)
            status, result = self.process_file(file_path)
            if status:
                success += 1
            else:
                failures.append((filename, result))

        print(f"\n上传完成！成功 {success} 个，失败 {len(failures)} 个")
        if failures:
            print("\n失败详情：")
            for filename, error in failures:
                print(f"{filename}: {error}")


if __name__ == "__main__":
    uploader = OpenWebUIBatchUploader()

    # 修改为您的文档目录路径
    doc_directory = "/home/aresen/下载/upload"

    if os.path.isdir(doc_directory):
        uploader.batch_upload(doc_directory)
    else:
        print(f"目录不存在: {doc_directory}")